#include "SplitFbxByMaterial.h"
#include <iostream>
#include <map>
#include <set>

vector<string> fbxTool::SplitFbxByMaterial::Split(const string& inputPath, const string& outputPath, int numParts)
{
	vector<string> result;

	vector<FbxNode*> meshNodes;
	FbxManager* manager; FbxScene* originalScene;
	GetAllMeshNodes(inputPath, meshNodes, manager, originalScene);
	// 按材质分组
	map<string, vector<FbxNode*>> materialGroups;
	for (FbxNode* node : meshNodes) {
		FbxMesh* mesh = (FbxMesh*)node->GetNodeAttribute();
		if (mesh) {
			int materialCount = node->GetMaterialCount();
			for (int i = 0; i < materialCount; ++i) {
				FbxSurfaceMaterial* material = node->GetMaterial(i);
				string materialName = material ? material->GetName() : "UnknownMaterial";
				materialGroups[materialName].push_back(node);
				break;
			}
		}
	}

	// 分割材质节点
	vector<vector<FbxNode*>> partitionedGroups(numParts);

	// 记录已分配的节点
	set<FbxNode*> assignedNodes;

	// 将材质节点分配到分组中
	for (const auto& group : materialGroups) {
		const string& materialName = group.first;
		const vector<FbxNode*>& nodes = group.second;

		for (FbxNode* node : nodes) {
			if (assignedNodes.count(node) == 0) {
				// 寻找最小的分组
				int bestGroup = 0;
				int minGroupSize = static_cast<int>(partitionedGroups[0].size());
				for (int i = 1; i < numParts; ++i) {
					if (partitionedGroups[i].size() < minGroupSize) {
						bestGroup = i;
						minGroupSize = static_cast<int>(partitionedGroups[i].size());
					}
				}
				partitionedGroups[bestGroup].push_back(node);
				assignedNodes.insert(node);
			}
		}
	}

	// 导出每个分组到独立的 FBX 文件
	for (int i = 0; i < numParts; ++i) {
		// 创建新的 FBX 管理器和设置
		FbxManager* partManager = FbxManager::Create();
		FbxIOSettings* partIosSettings = FbxIOSettings::Create(partManager, IOSROOT);
		partIosSettings->SetBoolProp(EXP_FBX_EMBEDDED, true);
		partManager->SetIOSettings(partIosSettings);

		// 创建新的场景和根节点
		FbxScene* newScene = FbxScene::Create(partManager, ("Scene_" + to_string(i + 1)).c_str());
		FbxNode* rootNode = FbxNode::Create(newScene, "RootNode");
		newScene->GetRootNode()->AddChild(rootNode);
		FbxGlobalSettings& globalSettings = newScene->GetGlobalSettings();
		globalSettings.SetAxisSystem(FbxAxisSystem::eOpenGL);

		// 克隆节点
		FbxCloneManager cloneManager;
		FbxCloneManager::CloneSet cloneSet;
		FbxCloneManager::CloneSetElement cloneOptions(FbxCloneManager::sConnectToClone, FbxCloneManager::sConnectToOriginal, FbxObject::eDeepClone);
		map<string, FbxAMatrix> transformMatrices;

		// 克隆每个分组中的节点
		for (FbxNode* node : partitionedGroups[i]) {
			FbxAMatrix globalTransform = node->EvaluateGlobalTransform();
			const string& nodeName = node->GetName();
			transformMatrices[nodeName] = globalTransform;
			cloneSet.Insert(node, cloneOptions);
			cloneManager.AddDependents(cloneSet, node, cloneOptions);
		}

		// 执行克隆
		if (!cloneManager.Clone(cloneSet, NULL)) {
			cerr << "克隆失败\n";
		}
		else {
			FbxCloneManager::CloneSet::RecordType* lIterator = cloneSet.Minimum();
			while (lIterator)
			{
				FbxObject* clonedObject = lIterator->GetValue().mObjectClone;  // 克隆对象

				if (FbxNode * clonedNode= FbxCast<FbxNode>(clonedObject)) {
					vector<FbxMesh*> meshes;  // 用于保存拆分后的网格
					const int materialCount = clonedNode->GetMaterialCount();
					FbxMesh* lMesh = clonedNode->GetMesh();
					for (int matIdx = 0; matIdx < materialCount; matIdx++)
					{
						// 创建一个新的 FbxMesh 用于存放该材质的面
						FbxMesh* newMesh = FbxMesh::Create(partManager, "SplitMesh");
						newMesh->InitControlPoints(lMesh->GetControlPointsCount());  // 初始化控制点

						// 通过材质拆分面
						FbxLayer* layer = lMesh->GetLayer(0);  // 获取网格的第一层
						FbxLayerElementMaterial* matElement = layer->GetMaterials();  // 获取材质元素

						if (matElement)
						{
							int polygonCount = lMesh->GetPolygonCount();  // 获取面的数量

							// 用于记录新网格的面
							std::vector<int> newPolygonIndices;
							int newPolygonIndex = 0;

							
							// 遍历所有面，将每个面分配到对应的材质网格中
							for (int polygonIdx = 0; polygonIdx < polygonCount; polygonIdx++)
							{
								int materialIndex = matElement->GetIndexArray().GetAt(polygonIdx);

								// 判断当前面是否属于当前材质
								if (materialIndex == matIdx)
								{
									int polygonSize = lMesh->GetPolygonSize(polygonIdx);
									newMesh->BeginPolygon();
									// 创建新的面并将原始网格的顶点索引复制到新网格
									for (int vertexIdx = 0; vertexIdx < polygonSize; vertexIdx++)
									{
										int controlPointIndex = lMesh->GetPolygonVertex(polygonIdx, vertexIdx);
										newMesh->AddPolygon(controlPointIndex);  // 向新网格中添加顶点
									}
									newMesh->EndPolygon();
									newPolygonIndices.push_back(newPolygonIndex++);
								}
							}
							
							// 将拆分后的网格添加到网格列表中
							meshes.push_back(newMesh);
						}
					}

					const string& nodeName = clonedNode->GetName();
					FbxAMatrix matrix = transformMatrices[nodeName];

					for (FbxMesh* mesh : meshes)
					{
						mesh->ConnectDstObject(newScene);
						// 创建一个新的 FbxNode，设置网格并添加到新节点
						FbxNode* newNode = FbxNode::Create(partManager, (nodeName + "_split").c_str());
						newNode->AddNodeAttribute(mesh);  // 将拆分后的网格添加到节点中

						// 设置新节点的变换
						newNode->LclTranslation.Set(matrix.GetT());  // 设置平移
						newNode->LclScaling.Set(matrix.GetS());      // 设置缩放
						newNode->LclRotation.Set(matrix.GetR());          // 设置旋转
						newNode->ConnectDstObject(newScene);
						// 将新节点添加到根节点
						rootNode->AddChild(newNode);
					}
					/*
					clonedNode->LclTranslation.Set(matrix.GetT());  // 设置平移
					clonedNode->LclScaling.Set(matrix.GetS());      // 设置缩放
					clonedNode->LclRotation.Set(matrix.GetR());
					rootNode->AddChild(clonedNode);  // 将克隆的节点添加到场景根节点
					*/
				}
				else {
					clonedObject->ConnectDstObject(newScene);
				}
				lIterator = lIterator->Successor();
			}
		}

		const string filename = ExportChildFbx(partitionedGroups, partManager, newScene, outputPath, i);
		if (!filename.empty())
			result.push_back(filename);
		
		partManager->Destroy();
	}

	manager->Destroy();
	return result;
}
